import torch
import torch.nn as nn
from torch.nn import Linear
import numpy as np

def smape_loss(y_pred, target):
    loss = 2 * (y_pred - target).abs() / (y_pred.abs() + target.abs() + 1e-8)
    return loss.mean()


# 生成decoder输入的mask矩阵（对角阵），看不到位置i以后的
def gen_trg_mask(length, device):
    # 对角矩阵
    mask = torch.tril(torch.ones(length, length, device=device)) == 1  # [len, len]
    # bsz, seq_len, feature_dim = trg.shape
    mask = (
        mask.float()
        .masked_fill(mask == 0, float("-inf"))
        .masked_fill(mask == 1, float(0.0))
    )

    return mask


# mask input的某个位置的输入
def gen_input_mask(keys: torch.tensor, queries_len, keys_len, head_num, device):
    # # keys [seq_len=keys_len, bsz, feature_dim]
    # # return [bsz*head_num, queries_len, keys_len]
    # bsz, seq_len, feature_dim = keys.shape
    # assert seq_len == keys_len, "seq_len != keys_len"
    # mask = np.arange(0, bsz*keys_len, dtype='float').reshape(bsz, keys_len)
    # for i in range(bsz):  # seq [keys_len, feature_dim]
    #     seq = keys[i]
    #     for j in range(keys_len):  # item [feature_dim]
    #         item = seq[j]
    #         if item[-1] == 0:  # 只要target 特征缺失，就不考虑，其他特征缺失也无关
    #             mask[i][j] = np.float('-inf')
    #         else:
    #             mask[i][j] = np.float(0)
    # mask = np.reshape(mask, (bsz, 1, keys_len))
    # mask = np.repeat(mask, head_num*queries_len, axis=1)  # [bsz, head_num*queries_len, keys_len]，每个query 考虑keys都是一样的
    # mask = np.reshape(mask, (bsz*head_num, queries_len, keys_len))  # [bsz*head_num, queries_len, keys_len]
    #
    # return torch.tensor(mask).to(device).float()
    # todo 不mask 缺失值，因为计算loss的时候就没计算缺失值
    return gen_generator_input_mask_(keys, queries_len, keys_len, head_num, device)


# mask input的某个位置的输入
def gen_generator_input_mask_(keys: torch.tensor, queries_len, keys_len, head_num, device):
    # keys [seq_len=keys_len, bsz, feature_dim]
    # return [bsz*head_num, queries_len, keys_len]
    bsz, seq_len, feature_dim = keys.shape
    assert seq_len == keys_len, "seq_len != keys_len"
    mask = np.zeros([bsz*head_num, queries_len, keys_len], dtype='float')

    return torch.tensor(mask).to(device).float()


class Encoder(nn.Module):
    def __init__(
        self,
        n_encoder_inputs,  # 输入的特征维度
        channels=512,  # 输出的特征维度
        dropout=0.1,
        lr=1e-4,
        num_layers=8,
        n_head=8
    ):
        super().__init__()

        # self.save_hyperparameters()

        self.lr = lr
        self.dropout = dropout
        self.n_head = n_head

        # 带参数的pos_embedding
        self.input_pos_embedding = torch.nn.Embedding(1024, embedding_dim=channels)  # input(*), output(*,channels)

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=channels,
            nhead=n_head,
            dropout=self.dropout,
            dim_feedforward=4 * channels,
        )

        self.encoder = torch.nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        # 将input_size -> channels
        self.input_projection = Linear(n_encoder_inputs, channels)
        self.do = nn.Dropout(p=self.dropout)
        self.sigmoid = nn.Sigmoid()

    def encode_src(self, src):
        # src: [bsz, seq_len, input_size]
        # return [seq_len, bsz, channel], [bsz*head_num, queries_len, keys_len]
        # [bsz, seq_len, input_size] -> [seq_len, bsz, channel]
        #  (seq, batch, feature). 默认batch_first=false
        src_start = self.input_projection(src).permute(1, 0, 2)

        in_sequence_len, batch_size = src_start.size(0), src_start.size(1)
        # [bsz*head_num, queries_len, keys_len]
        src_mask = gen_input_mask(src, in_sequence_len, in_sequence_len, head_num=self.n_head, device=src.device)

        pos_encoder = (
            torch.arange(0, in_sequence_len, device=src.device)  # [seq_len, ]
            .unsqueeze(0)  # [1, seq_len]
            .repeat(batch_size, 1)  # [bsz, seq_len]
        )
        # [bsz, seq_len] -> [bsz, seq_len, channel] -> [seq_len, bsz, channel]
        pos_encoder = self.input_pos_embedding(pos_encoder).permute(1, 0, 2)

        src = src_start + pos_encoder

        # 残差
        # [seq_len, bsz, channel]，mask就是src_mask，
        src = self.encoder(src, mask=src_mask) + src_start
        src = self.sigmoid(src)
        return src, src_mask

    # 输入input，原始数据，
    # return  encoder输出memory，input的mask矩阵用于memory的mask，
    def forward(self, x):
        src = x

        src, src_mask = self.encode_src(src)
        return src, src_mask


class Decoder(nn.Module):
    def __init__(
        self,
        n_decoder_inputs,  # decoder输入的特征维度
        n_decoder_outputs,  # decoder输出的特征维度，重构，则=encoder输入的特征维度
        channels=512,
        dropout=0.1,
        lr=1e-4,
        num_layers=8,
        n_head=8
    ):
        '''

        Args:
            n_decoder_inputs: decoder输入的维度
            n_decoder_outputs: decoder输出的维度，因为重构，n_decoder_inputs=n_decoder_outputs
        '''
        super().__init__()

        # self.save_hyperparameters()

        self.lr = lr
        self.dropout = dropout
        self.n_inputs = n_decoder_inputs
        self.n_outputs = n_decoder_outputs
        self.n_head = n_head

        # 带参数的pos_embedding
        self.target_pos_embedding = torch.nn.Embedding(1024, embedding_dim=channels)

        decoder_layer = nn.TransformerDecoderLayer(
            d_model=channels,
            nhead=n_head,
            dropout=self.dropout,
            dim_feedforward=4 * channels,
        )
        self.decoder = torch.nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        # 将input_size -> channels
        self.output_projection = Linear(n_decoder_inputs, channels)
        # decoder输出的维度，
        self.linear = Linear(channels, n_decoder_outputs)  #
        self.sigmoid = nn.Sigmoid()

    def decode_trg(self, trg, memory, memory_mask):
        '''
        teachder-forcing训练模式，decoder的输入是ground_truth
        如果是预测，则trg是上一步的输出
        :param trg: decoder的输入，[bsz, seq_len, feature_dim] // [seq_len, bsz, channel]
        :param memory: encoder的输出 [seq_len, bsz, channel]
        :param memory_mask:
            encoder的mask矩阵，由于decoder的作用是重构，memory与input一样mask同样的位置，因此复用memory_mask=src_mask
            [bsz*head_num, queries_len, keys_len]

        :return:
            [bsz, seq_len, output_dim]
        '''

        # [bsz, trg_seq_len, input_size] -> [seq_len, bsz, channel]
        # trg_start = self.output_projection(trg).permute(1, 0, 2)
        trg_start = memory  # todo 输入就是memory [seq_len, bsz, channel]

        out_sequence_len, batch_size = trg_start.size(0), trg_start.size(1)
        # 对角阵，看不到当前步后面的内容
        trg_mask = gen_trg_mask(out_sequence_len, memory.device)
        # decoder，输入也不看缺失值部分
        # trg_mask_2 = gen_input_mask(trg, out_sequence_len, out_sequence_len, self.n_head, trg.device)
        # trg_mask = trg_mask_2 + trg_mask  # 自动广播， [bsz*head_num, len, len]

        # memory_mask = gen_src_mask()

        pos_decoder = (
            torch.arange(0, out_sequence_len, device=memory.device)
            .unsqueeze(0)
            .repeat(batch_size, 1)
        )
        # [seq_len, bsz, embedding_dim=channel] -> [bsz, seq_len, embedding_dim]
        pos_decoder = self.target_pos_embedding(pos_decoder).permute(1, 0, 2)

        trg = pos_decoder + trg_start

        # [seq_len, bsz, channel]
        out = self.decoder(tgt=trg, memory=memory, tgt_mask=trg_mask, memory_mask=memory_mask) + trg_start

        # [seq_len, bsz, channel] -> [bsz, seq_len, channel]
        out = out.permute(1, 0, 2)
        # [bsz, seq_len, channel] -> [bsz, seq_len, output_dim]
        out = self.linear(out)
        out = self.sigmoid(out)
        return out

    def forward(self, memory, src_mask, x=None, free_running=False):
        '''
        Args:
            x: decoder的输入[seq_len, bsz, input_size]，训练teacher-forcing用真实数据， free_running模式则x是没有用的
            memory: encoder的输出，语义表示 [seq_len, bsz, channel]
            src_mask: encoder的mask矩阵，由于decoder的作用是重构，memory与input一样mask同样的位置，因此复用memory_mask=src_mask

        Returns:
            重构后的输出,[bsz, seq_len, output_dim]
        '''
        trg = x
        out = self.decode_trg(trg=trg, memory=memory, memory_mask=src_mask)
        return out


class Generator(nn.Module):
    def __init__(
        self,
        n_encoder_inputs,  # 输入特征的维度
        channels=512,  # 语义表示的维度，输出的特征维度
        dropout=0.1,
        lr=1e-4,
        num_layers=8,
        n_head=8
    ):
        super().__init__()

        # self.save_hyperparameters()

        self.lr = lr
        self.dropout = dropout
        # 带参数的pos_embedding
        self.input_pos_embedding = torch.nn.Embedding(1024, embedding_dim=channels)
        self.n_head = n_head

        encoder_layer = nn.TransformerEncoderLayer(
            d_model=channels,
            nhead=n_head,
            dropout=self.dropout,
            dim_feedforward=4 * channels,
        )

        self.encoder = torch.nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        # 将input_size -> channels
        self.input_projection = Linear(n_encoder_inputs, channels)
        self.do = nn.Dropout(p=self.dropout)
        self.sigmoid = nn.Sigmoid()

    def encode_src(self, src):
        # src: [bsz, seq_len, input_size]
        # [bsz, seq_len, input_size] -> [seq_len, bsz, channel]
        src_start = self.input_projection(src).permute(1, 0, 2)

        in_sequence_len, batch_size = src_start.size(0), src_start.size(1)
        src_mask = gen_generator_input_mask_(src, in_sequence_len, in_sequence_len, head_num=self.n_head, device=src.device)

        pos_encoder = (
            torch.arange(0, in_sequence_len, device=src.device)  # [seq_len, ]
                .unsqueeze(0)  # [1, seq_len]
                .repeat(batch_size, 1)  # [bsz, seq_len]
        )
        # [bsz, seq_len] -> [bsz, seq_len, channel] -> [seq_len, bsz, channel]
        pos_encoder = self.input_pos_embedding(pos_encoder).permute(1, 0, 2)

        src = src_start + pos_encoder

        # 残差
        # [seq_len, bsz, channel]
        src = self.encoder(src, mask=src_mask) + src_start
        src = self.sigmoid(src)
        return src, src_mask

    # 输入，随机向量，
    # 输出，memory，input_mask
    def forward(self, x):
        src = x

        src, src_mask = self.encode_src(src)

        return src, src_mask


class Discriminator(nn.Module):
    def __init__(
        self,
        n_decoder_inputs,  #
        n_decoder_outputs,
        channels=512,
        dropout=0.1,
        lr=1e-4,
        num_layers=8,
        n_head=8
    ):
        '''

        Args:
            n_decoder_inputs: decoder输入的维度, todo 暂时设顶为 feature_dim
            n_decoder_outputs: decoder输出的维度，todo 暂时设定为 feature_dim
            channels:
            dropout:
            lr:
        '''
        super().__init__()

        self.lr = lr
        self.dropout = dropout
        self.n_outputs = n_decoder_outputs
        self.n_inputs = n_decoder_inputs
        self.n_head = n_head
        # 带参数的pos_embedding
        self.target_pos_embedding = torch.nn.Embedding(1024, embedding_dim=channels)

        decoder_layer = nn.TransformerDecoderLayer(
            d_model=channels,
            nhead=n_head,
            dropout=self.dropout,
            dim_feedforward=4 * channels,
        )
        self.decoder = torch.nn.TransformerDecoder(decoder_layer, num_layers=num_layers)
        # 将input_size -> channels
        self.output_projection = Linear(n_decoder_inputs, channels)
        # decoder输出的维度，
        # self.linear = Linear(channels, n_decoder_outputs)  #
        # # 分类，输出0，1真实还是生成的数据，不要sigmoid，因为损失函数带了
        # self.classfiy = Linear(n_decoder_outputs, 1)
        self.linear = Linear(channels, 1)  # 由于不是free_running，直接linear输出0，1

    def decode_trg(self, memory_mask, trg, memory):
        '''
        teachder-forcing训练模式，decoder的输入是ground_truth
        如果是预测，则trg是上一步的输出
        :param trg: decoder的输入，[bsz, seq_len, feature_dim]
        :param memory: encoder的输出
        :param memory_mask: encoder的mask矩阵，decoder也mask input中mask的位置，因此memory_mask=src_mask

        :return:
        '''

        # [bsz, trg_seq_len, input_size] -> [seq_len, bsz, channel]
        # trg_start = self.output_projection(trg).permute(1, 0, 2)
        trg_start = memory  # decoder的输入就是memory [seq_len, bsz, channel]

        out_sequence_len, batch_size = trg_start.size(0), trg_start.size(1)
        # 对角阵，看不到当前步后面的内容
        trg_mask = gen_trg_mask(out_sequence_len, trg.device)
        # decoder，输入也不看缺失值部分
        # trg_mask_2 = gen_input_mask(trg, out_sequence_len, out_sequence_len, self.n_head, trg.device)
        # trg_mask = trg_mask_2 + trg_mask  # 自动广播， [bsz*head_num, len, len]

        # memory_mask = gen_src_mask()

        pos_decoder = (
            torch.arange(0, out_sequence_len, device=trg.device)
                .unsqueeze(0)
                .repeat(batch_size, 1)
        )
        # [seq_len, bsz, embedding_dim=channel] -> [bsz, seq_len, embedding_dim]
        pos_decoder = self.target_pos_embedding(pos_decoder).permute(1, 0, 2)

        trg = pos_decoder + trg_start

        out = self.decoder(tgt=trg, memory=memory, tgt_mask=trg_mask, memory_mask=memory_mask) + trg_start

        # [bsz, seq_len, channel] -> [seq_len, bsz, channel]
        out = out.permute(1, 0, 2)
        # [bsz, seq_len, channel] -> [bsz, seq_len, 1]
        out = self.linear(out)

        return out

    def forward(self, memory, src_mask):
        '''
        memory: [seq_len, bsz, channel]
        src_mask: encoder的mask矩阵，由于decoder的作用是重构，memory与input一样mask同样的位置，因此复用memory_mask=src_mask
        Returns:
        '''
        # todo 要不要mask掉缺失值
        return self.decode_trg(trg=memory, memory=memory, memory_mask=src_mask)


class Imputation(nn.Module):
    def __init__(self, opt, x):
        super(Imputation, self).__init__()

        self.z = torch.nn.Parameter(torch.tensor(x, dtype=torch.float32))

    def forward(self):
        return self.z

if __name__ == "__main__":
    pass

    # n_classes = 100
    # # [bsz, seq_len, feature_dim]
    # source = torch.rand(size=(32, 16, 9))
    # target_in = torch.rand(size=(32, 16, 9))
    # target_out = torch.rand(size=(32, 16, 1))
    #
    # ts = TimeSeriesForcasting(n_encoder_inputs=9, n_decoder_inputs=9)

